
import numpy as np

class DHAFLSystem:
    """DHA-FL核心框架"""
    def __init__(self, num_clients, model_dim):
        self.num_clients = num_clients
        self.model_dim = model_dim
        self.global_model = np.random.randn(model_dim)
        self.client_states = self._init_client_states()
        
    def _init_client_states(self):
        """初始化客户端状态"""
        return [{
            'model': np.random.randn(self.model_dim),
            'entropy_history': [],
            'hetero_score': np.random.uniform(0.5, 1.0),
            'layer': 'medium',
            'temperature': 0.8
        } for _ in range(self.num_clients)]
    
    def compute_heterogeneity_score(self, client_idx):
        """计算设备异构性评分 (式4)"""
        client = self.client_states[client_idx]
        # 简化的异构评分计算
        capability = np.exp(-0.8 * (1 - client['hetero_score']))
        comm_quality = 0.9 * client['hetero_score']
        divergence = np.random.uniform(0.1, 0.5)
        return 0.4*capability + 0.3*comm_quality + 0.3*divergence
    
    def compute_gradient_entropy(self, gradients):
        """梯度分布熵计算 (式5)"""
        # 简化的熵计算
        std = np.std(gradients)
        return np.log(std + 1e-8) if std > 0 else 0
    
    def update_temperature(self, entropy_change_rate):
        """温度系数动态演化 (式8)"""
        for client in self.client_states:
            if entropy_change_rate > 0.1:
                client['temperature'] *= 0.95  # 快速衰减
            else:
                client['temperature'] *= 0.98  # 缓慢回升
    
    def dynamic_hierarchization(self):
        """动态分层策略 (式10)"""
        for client in self.client_states:
            state_entropy = 0.6 * client['hetero_score'] + 0.4 * client['entropy_history'][-1]
            
            if state_entropy < 0.56:  # 高资源层
                client['layer'] = 'high'
            elif state_entropy < 0.7:  # 中资源层
                client['layer'] = 'medium'
            else:  # 低资源层
                client['layer'] = 'low'
    
    def attention_aggregation(self):
        """注意力聚合机制 (式6)"""
        aggregated_grad = np.zeros(self.model_dim)
        layer_weights = {'high': 0.5, 'medium': 0.33, 'low': 0.17}
        
        for layer in ['high', 'medium', 'low']:
            layer_clients = [c for c in self.client_states if c['layer'] == layer]
            if not layer_clients: continue
            
            entropies = [c['entropy_history'][-1] for c in layer_clients]
            weights = np.exp(-np.array(entropies) / np.mean([c['temperature'] for c in layer_clients]))
            weights /= np.sum(weights)
            
            layer_grad = np.sum([w * c['gradients'] for w, c in zip(weights, layer_clients)], axis=0)
            aggregated_grad += layer_weights[layer] * layer_grad
        
        return aggregated_grad
    
    def resource_adaptation(self, client_idx):
        """资源自适应 (3.4节)"""
        client = self.client_states[client_idx]
        if client['layer'] == 'low':
            # 轻量化模型部署
            pruning_mask = np.random.binomial(1, 0.7, size=self.model_dim)
            client['model'] *= pruning_mask
            
            # 通信压缩
            top_indices = np.argsort(np.abs(client['gradients']))[-int(self.model_dim*0.2):]
            return client['gradients'][top_indices], top_indices
        return client['gradients'], None
    
    def client_update(self, client_idx):
        """客户端本地更新"""
        client = self.client_states[client_idx]
        
        # 模拟梯度计算
        client['gradients'] = np.random.randn(self.model_dim)
        
        # 计算熵值并更新历史
        entropy = self.compute_gradient_entropy(client['gradients'])
        client['entropy_history'].append(entropy)
        
        # 资源自适应
        grad, _ = self.resource_adaptation(client_idx)
        return grad
    
    def server_aggregation(self):
        """服务器聚合"""
        # 更新温度系数
        entropy_changes = [np.mean(np.diff(c['entropy_history'][-5:])) for c in self.client_states]
        self.update_temperature(np.mean(entropy_changes))
        
        # 动态分层
        self.dynamic_hierarchization()
        
        # 注意力聚合
        aggregated_grad = self.attention_aggregation()
        
        # 更新全局模型
        self.global_model -= 0.01 * aggregated_grad
        return self.global_model

# 系统运行示例
if __name__ == "__main__":
    system = DHAFLSystem(num_clients=10, model_dim=100)
    
    for round_num in range(100):
        # 客户端更新
        for client_idx in range(10):
            system.client_update(client_idx)
        
        # 服务器聚合
        global_model = system.server_aggregation()
        print(f"Round {round_num}: Global model updated")
